/*
 * main.c
 *
 * Copyright (C) 2019 Alex A. <coderain@sdf.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stddef.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <assert.h>
#include <unistd.h>
#include <errno.h>
#include <linux/fs.h>
#include <linux/blkpg.h>
#include "nbd.h"

#define FUSE_USE_VERSION 26
#include <fuse.h>
#include <fuse_opt.h>

#define DEFAULT_HOST     "localhost"
#define DEFAULT_PORT     "10809"
#define FILE_NAME        "disk"
#define MAX_PARTITIONS    128

typedef struct
{
    char *tcp_host;
    char *tcp_port;
    char *unix_peer;
} fs_config_t;

typedef struct
{
    uint8_t type;
    uint64_t offset;
    uint64_t length;
} partition_t;

typedef struct
{
    fs_config_t config;
    nbd_conn_t conn;
    nbd_image_t image;
    time_t last_read_time, last_write_time;
    int owner_uid, owner_gid;
    uint8_t num_parts, max_part_num;
    partition_t partitions[MAX_PARTITIONS];
} fs_data_t;

enum
{
    NODE_NONEXISTENT = -1,
    NODE_ROOT_DIR = 0,
    NODE_WHOLE_FILE,
    NODE_PART_FILES,
};

static nbd_status_t rescan_mbr(fs_data_t *fs)
{
    uint8_t table[64];
    uint64_t extended = 0;
    fputs("[Rescanning partition table]\n", stderr);

    nbd_status_t error = nbd_read(&fs->conn, table, 0x1BE, sizeof(table));
    if (error) return error;

    fs->num_parts = fs->max_part_num = 0;
    memset(fs->partitions, 0, sizeof(fs->partitions));

    for (int i = 0; i < 4; i++)
    {
        uint8_t type = table[(i << 4) + 4];
        uint32_t start = le32toh(*(uint32_t*)&table[(i << 4) + 8]);
        uint32_t num_sectors = le32toh(*(uint32_t*)&table[(i << 4) + 12]);

        if (!type || type == 0x05 || type == 0x0F)
        {
            if (!extended && (type == 0x05 || type == 0x0F)) extended = (uint64_t)start << 9;
            continue;
        }

        fs->partitions[i].type   = type;
        fs->partitions[i].offset = (uint64_t)start << 9;
        fs->partitions[i].length = (uint64_t)num_sectors << 9;
        fs->num_parts++;
        fs->max_part_num = i + 1;
    }

    if (extended)
    {
        uint64_t current = 0;

        for (int i = 4; i < MAX_PARTITIONS; i++)
        {
            if (error = nbd_read(&fs->conn, table, extended + current + 0x1BE, sizeof(table) / 2)) return error;

            if (table[4])
            {
                fs->partitions[i].type   = table[4];
                fs->partitions[i].offset = (uint64_t)le32toh(*(uint32_t*)&table[8])  << 9;
                fs->partitions[i].length = (uint64_t)le32toh(*(uint32_t*)&table[12]) << 9;
                fs->num_parts++;
                fs->max_part_num = i + 1;
            }

            if (table[20] != 0x05 && table[20] != 0x0F) break;
            current = (uint64_t)le32toh(*(uint32_t*)&table[24]) << 9;
        }
    }

    return 0;
}

static int parse_path(fs_data_t *fs, const char *path)
{
    if (*path++ != '/') return NODE_NONEXISTENT;
    if (!*path) return NODE_ROOT_DIR;
    if (strncmp(path, FILE_NAME, sizeof(FILE_NAME) - 1) != 0) return NODE_NONEXISTENT;

    const char *str = &path[sizeof(FILE_NAME) - 1];
    if (!*str) return NODE_WHOLE_FILE;
    if (*str == '0') return NODE_NONEXISTENT;

    unsigned int number = 0;
    while (*str)
    {
        if (*str < '0' || *str > '9') return NODE_NONEXISTENT;
        number = number * 10 + *str++ - '0';
        if (number > fs->max_part_num) return NODE_NONEXISTENT;
    }

    return (number <= fs->max_part_num && fs->partitions[number - 1].type)
        ? NODE_PART_FILES + number - 1
        : NODE_NONEXISTENT;
}

static bool check_connection(fs_data_t *fs)
{
    if (nbd_is_connected(&fs->conn)) return true;

    if (fs->config.unix_peer
        ? nbd_connect_unix(fs->config.unix_peer, &fs->image, &fs->conn)
        : nbd_connect_tcp(fs->config.tcp_host, fs->config.tcp_port, &fs->image, &fs->conn))
    {
        rescan_mbr(fs);
        return true;
    }
    else
    {
        perror("Could not connect to NBD server");
        return false;
    }
}

static void *fs_init(struct fuse_conn_info *conn)
{
    struct fuse_context *ctx = fuse_get_context();
    fs_data_t *fs = ctx->private_data;

    fs->last_read_time = fs->last_write_time = time(NULL);
    fs->owner_uid = getuid();
    fs->owner_gid = getgid();
    nbd_init(&fs->conn);
    check_connection(fs);

    return fs;
}

static void fs_destroy(void *private_data)
{
    fs_data_t *fs = private_data;

    if (nbd_is_connected(&fs->conn))
    {
        nbd_flush(&fs->conn);
        nbd_disconnect(&fs->conn);
    }
}

static int fs_getattr(const char *path, struct stat *st)
{
    struct fuse_context *ctx = fuse_get_context();
    fs_data_t *fs = ctx->private_data;
    if (!check_connection(fs)) return -ENOTCONN;

    int node = parse_path(fs, path);

    switch (node)
    {
    case NODE_NONEXISTENT:
        return -ENOENT;

    case NODE_ROOT_DIR:
        st->st_mode = S_IFDIR | 0500;
        st->st_nlink = 3 + fs->num_parts;
        st->st_size = 512;
        break;

    default:
        st->st_mode = S_IFREG | 0600;
        st->st_nlink = 1;
        st->st_size = (node == NODE_WHOLE_FILE)
            ? fs->image.image_size
            : fs->partitions[node - NODE_PART_FILES].length;
        break;
    }

    st->st_ino = node + 1;
    st->st_uid = fs->owner_uid;
    st->st_gid = fs->owner_gid;
    st->st_atime = fs->last_read_time;
    st->st_mtime = st->st_ctime = fs->last_write_time;
    st->st_blksize = 512;
    st->st_blocks = (st->st_size + st->st_blksize - 1) / st->st_blksize;
    return 0;
}

static int fs_read(const char *path, char *buffer, size_t size, off_t _offset, struct fuse_file_info *fi)
{
    struct fuse_context *ctx = fuse_get_context();
    fs_data_t *fs = ctx->private_data;
    if (!check_connection(fs)) return -ENOTCONN;

    uint64_t offset = _offset;
    int node = parse_path(fs, path);

    switch (node)
    {
    case NODE_NONEXISTENT:
        return -ENOENT;

    case NODE_ROOT_DIR:
        return -EISDIR;

    default:
        /* Make sure the access is within the partition */
        if (offset >= fs->partitions[node - NODE_PART_FILES].length) return 0;

        if (size > fs->partitions[node - NODE_PART_FILES].length - offset)
        {
            size = fs->partitions[node - NODE_PART_FILES].length - offset;
        }

        offset += fs->partitions[node - NODE_PART_FILES].offset;
        /* Deliberately fall down */

    case NODE_WHOLE_FILE:
        /* Make sure the access is within the drive */
        if (offset >= fs->image.image_size) return 0;
        if (size > fs->image.image_size - offset) size = fs->image.image_size - offset;
        break;
    }

    if (!size) return 0;
    nbd_status_t status = nbd_read(&fs->conn, buffer, offset, size);
    return status ? -status : size;
}

static int fs_write(const char *path, const char *buffer, size_t size, off_t _offset, struct fuse_file_info *fi)
{
    struct fuse_context *ctx = fuse_get_context();
    fs_data_t *fs = ctx->private_data;
    if (!check_connection(fs)) return -ENOTCONN;

    uint64_t offset = _offset;
    int node = parse_path(fs, path);

    switch (node)
    {
    case NODE_NONEXISTENT:
        return -ENOENT;

    case NODE_ROOT_DIR:
        return -EISDIR;

    default:
        /* Make sure the access is within the partition */
        if (offset >= fs->partitions[node - NODE_PART_FILES].length) return -ENOSPC;

        if (size > fs->partitions[node - NODE_PART_FILES].length - offset)
        {
            size = fs->partitions[node - NODE_PART_FILES].length - offset;
        }

        offset += fs->partitions[node - NODE_PART_FILES].offset;
        /* Deliberately fall down */

    case NODE_WHOLE_FILE:
        /* Make sure the access is within the drive */
        if (offset >= fs->image.image_size) return -ENOSPC;

        if (size > fs->image.image_size - offset)
        {
            if (!(size = fs->image.image_size - offset)) return -ENOSPC;
        }

        break;
    }

    if (!size) return 0;
    nbd_status_t status = nbd_write(&fs->conn, buffer, offset, size);
    return status ? -status : size;
}

static int fs_readdir(const char *path, void *data, fuse_fill_dir_t filler, off_t off, struct fuse_file_info *fi)
{
    struct fuse_context *ctx = fuse_get_context();
    fs_data_t *fs = ctx->private_data;
    if (!check_connection(fs)) return -ENOTCONN;
    int node = parse_path(fs, path);

    switch (node)
    {
    case NODE_NONEXISTENT:
        return -ENOENT;

    case NODE_ROOT_DIR:
        filler(data, ".",  NULL, 0);
        filler(data, "..", NULL, 0);
        filler(data, FILE_NAME, NULL, 0);

        for (uint8_t i = 0; i < fs->max_part_num; i++)
        {
            if (fs->partitions[i].type)
            {
                int length = snprintf(NULL, 0, "%s%u", FILE_NAME, i + 1);
                char partname[length + 1];
                snprintf(partname, length + 1, "%s%u", FILE_NAME, i + 1);
                filler(data, partname, NULL, 0);
            }
        }

        return 0;

    default:
        return -ENOTDIR;
    }
}

static int fs_ioctl(const char *path, int cmd, void *arg, struct fuse_file_info *fi, unsigned int flags, void *data)
{
    struct fuse_context *ctx = fuse_get_context();
    fs_data_t *fs = ctx->private_data;
    if (!check_connection(fs)) return -ENOTCONN;

    int node = parse_path(fs, path);
    if (node == NODE_NONEXISTENT) return -ENOENT;
    if (node == NODE_ROOT_DIR) return -EISDIR;

    switch (cmd)
    {
    case BLKFLSBUF:
        return nbd_flush(&fs->conn);

    case BLKGETSIZE64:
        *(uint64_t*)data = (node == NODE_WHOLE_FILE)
            ? fs->image.image_size
            : fs->partitions[node - NODE_PART_FILES].length;
        return 0;

    case BLKRRPART:
        if (node != NODE_WHOLE_FILE) return -EINVAL;
        rescan_mbr(fs);
        return 0;

    default:
        return -ENOSYS;
    }
}

static int process_options(void *data, const char *arg, int key, struct fuse_args *args)
{
    fs_config_t *config = data;
}

int main(int argc, char *argv[])
{
    static const struct fuse_opt options[] = {
        { "--tcp=%s",  offsetof(fs_config_t, tcp_host),  0 },
        { "-t %s",     offsetof(fs_config_t, tcp_host),  0 },
        { "--port=%s", offsetof(fs_config_t, tcp_port),  0 },
        { "-p %s",     offsetof(fs_config_t, tcp_port),  0 },
        { "--unix=%s", offsetof(fs_config_t, unix_peer), 0 },
        { "-u %s",     offsetof(fs_config_t, unix_peer), 0 },
        FUSE_OPT_END
    };

    struct fuse_args args = FUSE_ARGS_INIT(argc, argv);
    static fs_data_t data = { .config = { .tcp_host = DEFAULT_HOST, .tcp_port = DEFAULT_PORT, .unix_peer = NULL } };
    fuse_opt_parse(&args, &data.config, options, &process_options);

    /*
     * We must add this or write() performance will be absolutely horrible.
     * No I'm not kidding. It gets slower than dial-up networking!
     *
     * TODO: Implement internal caching so we're not at the mercy of buffer sizes.
     */
    fuse_opt_add_arg(&args, "-obig_writes");

    static const struct fuse_operations operations = {
        .init     = fs_init,
        .destroy  = fs_destroy,
        .getattr  = fs_getattr,
        .read     = fs_read,
        .write    = fs_write,
        .readdir  = fs_readdir,
        .ioctl    = fs_ioctl,
    };

    return fuse_main(args.argc, args.argv, &operations, &data);
}
